import matplotlib
matplotlib.use('TkAgg')
from matplotlib import rcParams
rcParams['font.sans-serif'] = ['SimHei']  # 中文字体
rcParams['axes.unicode_minus'] = False    # 正常显示负号
import pandas as pd
import numpy as np
import joblib
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import (
    classification_report, accuracy_score,
    f1_score, precision_score, recall_score,
    confusion_matrix
)
from word2vec_processor import Word2VecProcessor
from utils.config import Config

# ================= 配置路径 =================
root_path = "D:/pycode/group4_nlp_project"
conf = Config(root_path)
data_path = conf.train_path
model_path = conf.basic_model_path
cat_rf_model_path = model_path + "cat_rf.pkl"
cat_w2v_path = model_path + "cat_word2vec.pkl"


def train_cat_model(data_path=data_path):
    # 1. 读取数据
    df = pd.read_csv(data_path)

    # 2. 训练 Word2Vec
    print("训练 Word2Vec 模型...")
    w2v = Word2VecProcessor(vector_size=100, window=5, min_count=3, epochs=10)
    w2v.train(df, text_column="review_clean")
    w2v.save(cat_w2v_path)

    # 3. 转换句子向量
    print("生成句子向量...")
    df["vector"] = df["review_clean"].apply(lambda x: w2v.get_sentence_vector(str(x)))
    df = df[df["vector"].notnull()].reset_index(drop=True)
    X = np.vstack(df["vector"].values)
    y_cat = df["cat"]

    # 4. 切分数据集
    X_train, X_test, y_train, y_test = train_test_split(
        X, y_cat, test_size=0.2, random_state=42, stratify=y_cat
    )

    # 5. 训练随机森林
    print("训练随机森林分类器...")
    rf = RandomForestClassifier(n_estimators=200, random_state=917)
    rf.fit(X_train, y_train)

    # 6. 预测与评估
    y_pred = rf.predict(X_test)
    print("\n==== 多分类任务 (随机森林 - cat) ====")
    print(classification_report(y_test, y_pred))

    # 计算核心指标
    acc = accuracy_score(y_test, y_pred)
    f1_macro = f1_score(y_test, y_pred, average="macro")
    f1_micro = f1_score(y_test, y_pred, average="micro")
    prec_macro = precision_score(y_test, y_pred, average="macro")
    rec_macro = recall_score(y_test, y_pred, average="macro")

    # 输出指标
    print("整体准确率:", acc)
    print("宏平均F1:", f1_macro)
    print("微平均F1:", f1_micro)
    print("宏平均Precision:", prec_macro)
    print("宏平均Recall:", rec_macro)

    # === 输出指标表格 ===
    metrics_df = pd.DataFrame({
        "Accuracy": [acc],
        "Macro-F1": [f1_macro],
        "Micro-F1": [f1_micro],
        "Macro-Precision": [prec_macro],
        "Macro-Recall": [rec_macro]
    })
    print("\n==== 指标汇总 ====")
    print(metrics_df.to_string(index=False))

    # === 绘制混淆矩阵 ===
    cm = confusion_matrix(y_test, y_pred, labels=rf.classes_)
    plt.figure(figsize=(8, 6))
    sns.heatmap(cm, annot=True, fmt="d", cmap="Blues",
                xticklabels=rf.classes_, yticklabels=rf.classes_)
    plt.xlabel("预测值")
    plt.ylabel("真实值")
    plt.title("混淆矩阵 (随机森林 - cat)")
    plt.tight_layout()
    plt.show()

    # 7. 保存模型和 Word2Vec
    joblib.dump(rf, cat_rf_model_path)
    print(f"\n模型已保存: {cat_rf_model_path}")
    print(f"词向量已保存: {cat_w2v_path}")


if __name__ == "__main__":
    train_cat_model()
